package com.itranswarp.learnjava.utils;

import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import org.apache.tomcat.util.http.fileupload.ByteArrayOutputStream;

import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

public class CachedHttpServletResponse extends HttpServletResponseWrapper {

    private ByteArrayOutputStream output = new ByteArrayOutputStream();
    private PrintWriter writer;
    private ServletOutputStream outputStream;
    private boolean writerUsed = false;
    private boolean streamUsed = false;

    public CachedHttpServletResponse(HttpServletResponse response) {
        super(response);
    }

    @Override
    public PrintWriter getWriter() throws IOException {
        if (streamUsed) {
            throw new IllegalStateException("getOutputStream() has already been called!");
        }
        if (writer == null) {
            writerUsed = true;
            writer = new PrintWriter(new OutputStreamWriter(output, StandardCharsets.UTF_8), true);
        }
        return writer;
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        if (writerUsed) {
            throw new IllegalStateException("getWriter() has already been called!");
        }
        if (outputStream == null) {
            streamUsed = true;
            outputStream = new ServletOutputStream() {
                @Override
                public boolean isReady() {
                    return true;
                }

                @Override
                public void setWriteListener(WriteListener listener) {
                }

                @Override
                public void write(int b) throws IOException {
                    output.write(b);
                }
            };
        }
        return outputStream;
    }

    public byte[] getContent() {
        if (writer != null) {
            writer.flush(); // 确保写入
        }
        return output.toByteArray();
    }
}
